#ifndef __XDO_02__
#define __XDO_02__

#include "../util/utils.pe.h"

namespace xm
{
	class XdoXPackB : public IKDetector, UnkBase
	{
	protected:

		struct DasmCallback : IPEDisAsmCallback
		{
			STDMETHOD(OnEvent)( DASM_TRACE_EVENT uEvent )
			{
				return E_FAIL;
			}
			STDMETHOD(OnInstruction)( X64INST * lpInst )
			{
				return E_FAIL;
			}
		};

	public:
		UNKNOWN_IMP1(IKDetector);
		DEFAULT_INIT_CLASS();
		STDMETHOD(Detect)( IKDetectSession * sess, DetectResult * pdr )
		{
			xm::iptr<IPeReader> winpe( sess->CurrentArchive() );
			if( !winpe ) return E_FAIL;
			if( winpe->IsWow64() ) return E_FAIL;
			
			DWORD dwMask = 0;
			RFAILED( TestInfectOrPacked(winpe,dwMask) );
			if( dwMask ) return E_FAIL;
			
			DWORD dwEntry = 0;
			if( !winpe->EntryPointRva(&dwEntry) ) 
				return E_FAIL;
			
			iptr<IKPETools> dasm( PROGID_WinPETools );
			if( !dasm ) return E_FAIL;
			
			RFAILED( dasm->Init( winpe ) );
			DasmCallback cb;
			dasm->FetchBIL( dwEntry, DTO_CALL, &cb );
			
			return E_FAIL;
		}
		STDMETHOD(Treat)( IKDetectSession * ctx, DetectResult * pdr )
		{
			return ctx->RawStream()->Remove();
		}
		STDMETHOD(Reset)()
		{
			return S_OK;
		}
	};

	class XdoXPackC : public IKDetector, UnkBase
	{
	public:
		UNKNOWN_IMP1(IKDetector);
		DEFAULT_INIT_CLASS();
		STDMETHOD(Detect)( IKDetectSession * sess, DetectResult * pdr )
		{
			xm::iptr<IPeReader> winpe( sess->CurrentArchive() );
			if( !winpe ) return E_FAIL;
			if( winpe->IsWow64() ) return E_FAIL;

			WinPeParser::TImgFileHdr * fh = winpe->FileHeader();
			if( !fh ) return E_FAIL;

			if( fh->Characteristics & WINPE_IMAGE_FILE_DLL ) 
				return E_FAIL;
			
			DWORD dwMask = 0;
			RFAILED( TestInfectOrPacked(winpe,dwMask) );
			if( dwMask ) return E_FAIL;
			
			DWORD num_sect = 0;
			RASSERT( winpe->SectionCount(&num_sect), E_FAIL );

			DWORD cbImage = 0;
			RASSERT( winpe->ImageBaseVa( NULL, &cbImage ), E_FAIL );

			DOUBLE cbBSS = 0;
			DOUBLE cBSSSections = 0;
			for( DWORD i = 0; i < num_sect; ++ i )
			{
				IPeReader::TImgSecHdr * hdr = winpe->SectionHeader(i);
				if( !tcs::nicmp( (CHAR*)hdr->Name, ".ndata", 7 ) )
					continue;

				if( hdr->SizeOfRawData == 0 )
				{
					if( !tcs::ncmp( (CHAR*)hdr->Name, "BSS", 4 ) )
						continue;

					cbBSS += hdr->Misc.VirtualSize;
					cBSSSections ++;
				}
				else if( hdr->SizeOfRawData > 0x1000 )
				{
					DOUBLE time = (DOUBLE)hdr->Misc.VirtualSize / (DOUBLE)hdr->SizeOfRawData;
					if( time > 3000 )
					{
						return S_OK;
					}
				}
			}

			if( cBSSSections/(DOUBLE)num_sect < 0.20 )
				return E_FAIL;

			if( (DOUBLE)cbImage - cbBSS < 32*1024 )
				return E_FAIL;

			if( cbBSS < 1024 * 1024 ) 
				return E_FAIL;
			
			if( cbBSS / (DOUBLE)cbImage > 0.8 ) 
				return S_OK;
			
			return E_FAIL;
		}
		STDMETHOD(Treat)( IKDetectSession * ctx, DetectResult * pdr )
		{
			return ctx->RawStream()->Remove();
		}
		STDMETHOD(Reset)()
		{
			return S_OK;
		}
	};

	class XdoAntiEmuA: public IKDetector, UnkBase
	{
	protected:

		typedef std::vector<X64INSTEX> InstVector;

		struct DasmCallback : IPEDisAsmCallback
		{
			InstVector	InstList;

			STDMETHOD(OnEvent)( DASM_TRACE_EVENT uEvent )
			{
				return S_OK;
			}
			STDMETHOD(OnInstruction)( X64INST * lpInst )
			{
				InstList.push_back( *(X64INSTEX*)lpInst );
				return InstList.size() < 32 ? S_OK : E_FAIL;
			}
		};

	protected:

		bool is_malware_api( uint32_t iat_va, IKPETools * tool )
		{
			IKPETools::IMPITEM * impitem = tool->ImpFindByIAT( iat_va, TRUE );
			if( !impitem ) return false;
			if( !strcmp( impitem->api, "GetModuleHandleA" ) ||
				!strcmp( impitem->api, "GetModuleHandleW" ) ||
				!strcmp( impitem->api, "LoadLibraryA" ) ||
				!strcmp( impitem->api, "LoadLibraryW" ) )
			{
				return true;
			}
			return false;
		}

	public:
		UNKNOWN_IMP1(IKDetector);
		DEFAULT_INIT_CLASS();

		STDMETHOD(Detect)( IKDetectSession * sess, DetectResult * pdr )
		{
			xm::iptr<IPeReader> winpe( sess->CurrentArchive() );
			if( !winpe ) return E_FAIL;
			if( winpe->IsWow64() ) return E_FAIL;

			DWORD dwMask = 0;
			RFAILED( TestInfectOrPacked(winpe,dwMask) );
			if( dwMask ) return E_FAIL;

			DWORD dwEntry = 0;
			if( !winpe->EntryPointRva(&dwEntry) ) 
				return E_FAIL;

			iptr<IKPETools> dasm;
			RFAILED( get_this_pe_tools( sess, winpe, dasm.pp() ) );

			DasmCallback cb;
			dasm->FetchBIL( dwEntry, DTO_CALL|DTO_BRACH_FIRST|DTO_RVA_ONCE|DTO_EX_INS, &cb );

			size_t c = cb.InstList.size();
			if( !c ) return E_FAIL;

			uint8_t opbuf[32] = {};
			size_t  opbs = 0;

			for( size_t i = 0; i < c; ++ i )
			{
				X64INST & inst = cb.InstList[i];
				if( inst.OpCode != 0xE8 )
					continue;

				if( i + 3 >= c ) return E_FAIL;

				X64INSTEX & ki = cb.InstList[i+1];
				if( ki.OpCode == 0xA1 )
				{
					if( cb.InstList[i+2].OpCode != 0x50 || cb.InstList[i+3].OpCode != 0xC3 )
						return E_FAIL;

					uint32_t iat_va = ki.mix.i32.imm.imm32;
					if( !is_malware_api( iat_va, dasm ) )
						return E_FAIL;
					return S_OK;
				}
				else if( ki.OpCode == 0xFF35 )
				{
					if( cb.InstList[i+2].OpCode != 0xC3 )
						return E_FAIL;

					uint32_t iat_va = ki.mix.i32.disp.disp32;
					if( !is_malware_api( iat_va, dasm ) )
						return E_FAIL;
					return S_OK;
				}
				else
				{
					return E_FAIL;
				}
			}

			return E_FAIL;
		}
		STDMETHOD(Treat)( IKDetectSession * ctx, DetectResult * pdr )
		{
			return ctx->RawStream()->Remove();
		}
		STDMETHOD(Reset)()
		{
			return S_OK;
		}
	};

};


#endif